import pandas as pd
import numpy as np
from sklearn.preprocessing import StandardScaler
import tensorflow as tf

class TKJ(object):
    def __init__(self):
        self.df = pd.read_csv("regression_b/scenic_data.csv")
        
    def create_dataset(self, data, n_steps):
        x, y = [], []
        for i in range(len(data) - n_steps):
            x_values = data[i : i+n_steps]
            x_values = pd.DataFrame(x_values)
            x_values.iloc[-1, x_values.columns.get_loc('count')] = 0
            x.append(x_values.values)
            y.append(data[i+n_steps, 0])
        return np.array(x), np.array(y)
    
    def get_model(self):
        n_steps = 7  # 定义n_steps
        data = self.df.values
        x, y = self.create_dataset(data, n_steps)
        
        # 划分训练集与验证集
        train_size = int(len(x) * 0.8)
        x_train, x_test = x[:train_size], x[train_size:]
        y_train, y_test = y[:train_size], y[train_size:]
        
        # 训练模型
        model = tf.keras.models.Sequential()
        model.add(tf.keras.layers.LSTM(64, activation='relu', 
        return_sequences=True, input_shape=(n_steps, x_train.shape[2])))
        model.add(tf.keras.layers.LSTM(64, activation='relu'))
        model.add(tf.keras.layers.Dense(1))
        model.compile(optimizer='adam', loss='mse')
        model.fit(x_train, y_train, epochs=50, validation_data=(x_test, y_test))
        
        # 评估
        loss = model.evaluate(x_test, y_test)
        print(f"模型测试损失：{loss}")
        
        # 保存模型
        tf.keras.models.save_model(model, "regression_b/my_model.keras")

if __name__ == '__main__':
    nn = ML()
    nn.get_model()
    